package org.jetlinks.community.datasource;

import lombok.AllArgsConstructor;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;

import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

@Slf4j
public class DefaultDataSourceManager implements DataSourceManager {

    private final Map<String, DataSourceProvider> providers = new ConcurrentHashMap<>();

    private final Map<CacheKey, DataSourceRef> cachedDataSources = new ConcurrentHashMap<>();

    private final DataSourceConfigManager dataSourceConfigManager;

    public DefaultDataSourceManager(DataSourceConfigManager configManager) {
        this.dataSourceConfigManager = configManager;
        this.dataSourceConfigManager
            .doOnConfigChanged((state, properties) -> {
                if (state == DataSourceConfigManager.ConfigState.disabled) {
                    this.removeDataSource(properties.getTypeId(), properties.getId());
                } else {
                    this
                        .reloadDataSource(properties.getTypeId(), properties.getId())
                        .subscribe();
                }

            });
    }

    public void register(DataSourceProvider provider) {
        log.debug("Register DataSource {} Provider {}", provider.getType().getId(), provider);
        providers.put(provider.getType().getId(), provider);
    }

    public void register(DataSource dataSource) {
        log.debug("Register DataSource {} {}", dataSource.getType().getId(), dataSource);
        CacheKey key = new CacheKey(dataSource.getType().getId(), dataSource.getId());
        DataSourceRef ref = new DataSourceRef(Mono.just(dataSource));
        cachedDataSources.put(key, ref);
    }

    @Override
    public Flux<DataSource> getDataSources(String typeId) {
        return Flux
            .fromIterable(cachedDataSources.values())
            .flatMap(DataSourceRef::getRef)
            .filter(dataSource -> Objects.equals(typeId, dataSource.getType().getId()));
    }

    @Override
    public List<DataSourceType> getSupportedType() {
        return providers
            .values()
            .stream()
            .map(DataSourceProvider::getType)
            .collect(Collectors.toList());
    }

    @Override
    public Mono<DataSource> getDataSource(DataSourceType type,
                                          String datasourceId) {
        return getDataSource(type.getId(), datasourceId);
    }

    @Override
    public Mono<DataSource> getDataSource(String typeId, String datasourceId) {
        return getOrCreateRef(typeId, datasourceId).getRef();
    }

    private DataSourceRef getOrCreateRef(String typeId, String datasourceId) {
        return cachedDataSources
            .computeIfAbsent(new CacheKey(typeId, datasourceId),
                             key -> new DataSourceRef(loadDataSource(key.type, key.datasourceId).cache()));

    }

    public Mono<Tuple2<DataSourceConfig, DataSourceProvider>> loadConfigAndProvider(String typeId, String datasourceId) {

        return Mono
            .zip(
                dataSourceConfigManager
                    .getConfig(typeId, datasourceId),
                Mono
                    .justOrEmpty(providers.get(typeId))
                    .switchIfEmpty(Mono.error(() -> new UnsupportedOperationException("unsupported datasource type " + typeId)))
            );

    }

    public Mono<DataSource> loadDataSource(String typeId, String datasourceId) {

        return loadConfigAndProvider(typeId, datasourceId)
            .flatMap(tp2 -> tp2.getT2().createDataSource(tp2.getT1()))
            .doOnNext(dataSource -> log.debug("load {} datasource [{}]", dataSource
                .getType()
                .getId(), dataSource.getId()));

    }

    private void removeDataSource(String typeId, String datasourceId) {
        DataSourceRef cache = cachedDataSources.remove(new CacheKey(typeId, datasourceId));
        if (cache != null) {
            cache.dispose();
            log.debug("removed {} datasource [{}]", typeId, datasourceId);
        }
    }

    private Mono<DataSource> reloadDataSource(String typeId, String datasourceId) {
        DataSourceRef cache = getOrCreateRef(typeId, datasourceId);

        return loadConfigAndProvider(typeId, datasourceId)
            .flatMap(tp2 -> cache.getRef()
                                 .flatMap(datasource -> tp2.getT2().reload(datasource, tp2.getT1()))
                                 .doOnNext(datasource -> {
                                     log.debug("reload {} datasource [{}]", datasource
                                         .getType()
                                         .getId(), datasource.getId());
                                     cache.ref = Mono.just(datasource);
                                 })
            )
            .doOnError(err -> log.debug("reload {} datasource [{}] error ", typeId, datasourceId, err));

    }

    @AllArgsConstructor
    @EqualsAndHashCode
    static class CacheKey {
        private final String type;
        private final String datasourceId;
    }

    static class DataSourceRef implements Disposable {

        @Getter
        private volatile Mono<DataSource> ref;

        private boolean disposed = false;

        public DataSourceRef(Mono<DataSource> ref) {
            this.ref = Mono.defer(() -> ref);
        }

        @Override
        public void dispose() {
            ref = Mono.empty();
            disposed = true;
        }

        @Override
        public boolean isDisposed() {
            return disposed;
        }
    }

}
